#include <iostream>
#include <string>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <pcap.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <netinet/udp.h>
#include <netinet/ether.h>
#include <netinet/ip_icmp.h>
#include <arpa/inet.h>
#include <unordered_map>
#include <vector>
#include <algorithm>
#include <ctime>
#include <chrono>
#include <sstream>

// 协议类型定义
enum Protocol {
    UNKNOWN,
    TCP,
    UDP,
    ICMP
};

// 攻击类型定义
enum AttackType {
    NONE,
    SYN_FLOOD,
    UDP_FLOOD,
    ICMP_FLOOD,
    HTTP_FLOOD,
    SLOWLORIS,
    FRAGMENTATION_ATTACK
};

// 会话状态结构体
struct Session {
    std::string src_ip;
    std::string dst_ip;
    uint16_t src_port;
    uint16_t dst_port;
    Protocol protocol;
    time_t first_seen;
    time_t last_seen;
    int packet_count;
    int total_bytes;
    
    // TCP特定状态
    uint32_t last_seq;
    uint32_t last_ack;
    bool is_syn;
    bool is_ack;
    bool is_fin;
    int syn_count;
    int ack_count;
    int fin_count;
    
    // 流量统计
    std::unordered_map<time_t, int> packet_rate; // 按秒统计的包数量
};

// 规则结构体
struct Rule {
    std::string name;
    std::string description;
    std::string src_ip;    // 源IP，*表示任意
    std::string dst_ip;    // 目的IP，*表示任意
    uint16_t src_port;     // 源端口，0表示任意
    uint16_t dst_port;     // 目的端口，0表示任意
    Protocol protocol;     // 协议类型，0表示任意
    std::string payload_pattern;  // payload匹配模式
    bool alert_on_match;   // 匹配时是否告警
    int threshold;         // 阈值，根据规则类型有不同含义
};

// 目标统计结构体 - 用于检测针对同一目标的攻击
struct TargetStats {
    std::string ip;
    time_t first_seen;
    std::unordered_map<std::string, int> src_ips; // 来源IP及包数量
    std::unordered_map<Protocol, int> proto_counts; // 协议分布
    std::unordered_map<time_t, int> packet_rate; // 按秒统计的包数量
    std::unordered_map<time_t, int> byte_rate; // 按秒统计的字节数
    int total_packets;
    int total_bytes;
};

// 全局变量
std::unordered_map<std::string, Session> sessions;
std::unordered_map<std::string, TargetStats> target_stats; // 按目标IP统计
std::vector<Rule> detection_rules;
int packet_counter = 0;
time_t program_start_time;

// 配置参数 - 可根据需要调整
const int SYN_FLOOD_THRESHOLD = 50; // 每秒SYN包阈值
const int UDP_FLOOD_THRESHOLD = 100; // 每秒UDP包阈值
const int ICMP_FLOOD_THRESHOLD = 30; // 每秒ICMP包阈值
const int HTTP_FLOOD_THRESHOLD = 60; // 每秒HTTP请求阈值
const int FRAGMENT_THRESHOLD = 20; // 同一IP的分片包阈值
const int LARGE_UDP_SIZE = 1024; // 大UDP包大小阈值
const int SLOWLORIS_TIMEOUT = 30; // Slowloris连接超时阈值(秒)
const int SLOWLORIS_THRESHOLD = 10; // 来自同一IP的未完成HTTP连接阈值

// 生成会话ID
std::string generate_session_id(const std::string& src_ip, const std::string& dst_ip, 
                               uint16_t src_port, uint16_t dst_port, Protocol protocol) {
    if (src_ip < dst_ip || (src_ip == dst_ip && src_port < dst_port)) {
        return src_ip + ":" + std::to_string(src_port) + "-" + 
               dst_ip + ":" + std::to_string(dst_port) + "-" + std::to_string(protocol);
    } else {
        return dst_ip + ":" + std::to_string(dst_port) + "-" + 
               src_ip + ":" + std::to_string(src_port) + "-" + std::to_string(protocol);
    }
}

// 获取映射表中的值，处理键不存在的情况
template<typename K, typename V>
V get_map_value(const std::unordered_map<K, V>& map, const K& key, const V& default_val = V()) {
    auto it = map.find(key);
    return (it != map.end()) ? it->second : default_val;
}

// 更新目标统计信息
void update_target_stats(const std::string& src_ip, const std::string& dst_ip,
                        Protocol protocol, int packet_size) {
    time_t now = time(nullptr);
    
    // 初始化或更新目标统计
    if (target_stats.find(dst_ip) == target_stats.end()) {
        TargetStats new_stats;
        new_stats.ip = dst_ip;
        new_stats.first_seen = now;
        new_stats.total_packets = 0;
        new_stats.total_bytes = 0;
        target_stats[dst_ip] = new_stats;
    }
    
    TargetStats& stats = target_stats[dst_ip];
    stats.total_packets++;
    stats.total_bytes += packet_size;
    
    // 更新来源IP统计
    stats.src_ips[src_ip]++;
    
    // 更新协议统计
    stats.proto_counts[protocol]++;
    
    // 更新速率统计（按秒）
    stats.packet_rate[now]++;
    stats.byte_rate[now] += packet_size;
    
    // 清理过期数据（只保留最近60秒）
    time_t cutoff = now - 60;
    for (auto it = stats.packet_rate.begin(); it != stats.packet_rate.end();) {
        if (it->first < cutoff) {
            it = stats.packet_rate.erase(it);
        } else {
            ++it;
        }
    }
    
    for (auto it = stats.byte_rate.begin(); it != stats.byte_rate.end();) {
        if (it->first < cutoff) {
            it = stats.byte_rate.erase(it);
        } else {
            ++it;
        }
    }
}

// 更新会话状态
void update_session(const std::string& src_ip, const std::string& dst_ip,
                   uint16_t src_port, uint16_t dst_port, Protocol protocol,
                   const u_char* payload __attribute__((unused)), int payload_len,
                   const tcphdr* tcp_header = nullptr) {
    std::string session_id = generate_session_id(src_ip, dst_ip, src_port, dst_port, protocol);
    time_t now = time(nullptr);
    int packet_size = payload_len + (tcp_header ? (tcp_header->doff * 4) : 0) + (protocol == TCP ? sizeof(tcphdr) : 0);
    
    if (sessions.find(session_id) == sessions.end()) {
        Session new_session;
        new_session.src_ip = src_ip;
        new_session.dst_ip = dst_ip;
        new_session.src_port = src_port;
        new_session.dst_port = dst_port;
        new_session.protocol = protocol;
        new_session.first_seen = now;
        new_session.last_seen = now;
        new_session.packet_count = 0;
        new_session.total_bytes = 0;
        new_session.last_seq = 0;
        new_session.last_ack = 0;
        new_session.is_syn = false;
        new_session.is_ack = false;
        new_session.is_fin = false;
        new_session.syn_count = 0;
        new_session.ack_count = 0;
        new_session.fin_count = 0;
        sessions[session_id] = new_session;
    }
    
    Session& session = sessions[session_id];
    session.last_seen = now;
    session.packet_count++;
    session.total_bytes += packet_size;
    session.packet_rate[now]++;
    
    // 清理会话中过期的速率数据
    time_t cutoff = now - 60;
    for (auto it = session.packet_rate.begin(); it != session.packet_rate.end();) {
        if (it->first < cutoff) {
            it = session.packet_rate.erase(it);
        } else {
            ++it;
        }
    }
    
    if (protocol == TCP && tcp_header != nullptr) {
        session.last_seq = ntohl(tcp_header->seq);
        session.last_ack = ntohl(tcp_header->ack_seq);
        
        if (tcp_header->syn == 1) {
            session.is_syn = true;
            session.syn_count++;
        }
        if (tcp_header->ack == 1) {
            session.is_ack = true;
            session.ack_count++;
        }
        if (tcp_header->fin == 1) {
            session.is_fin = true;
            session.fin_count++;
        }
    }
}

// 加载检测规则
void load_rules() {
    // SSH暴力破解检测规则
    Rule ssh_brute_force;
    ssh_brute_force.name = "SSH_BRUTE_FORCE";
    ssh_brute_force.description = "检测SSH暴力破解尝试";
    ssh_brute_force.src_ip = "*";
    ssh_brute_force.dst_ip = "*";
    ssh_brute_force.src_port = 0;
    ssh_brute_force.dst_port = 22;
    ssh_brute_force.protocol = TCP;
    ssh_brute_force.payload_pattern = "";
    ssh_brute_force.alert_on_match = true;
    ssh_brute_force.threshold = 5; // 超过5个包则告警
    detection_rules.push_back(ssh_brute_force);
    
    // SQL注入检测规则
    Rule sql_injection;
    sql_injection.name = "SQL_INJECTION";
    sql_injection.description = "检测HTTP请求中的SQL注入模式";
    sql_injection.src_ip = "*";
    sql_injection.dst_ip = "*";
    sql_injection.src_port = 0;
    sql_injection.dst_port = 80;
    sql_injection.protocol = TCP;
    sql_injection.payload_pattern = "UNION SELECT|OR 1=1|' OR '1'='1|INSERT INTO|DELETE FROM|DROP TABLE";
    sql_injection.alert_on_match = true;
    sql_injection.threshold = 1; // 只要匹配就告警
    detection_rules.push_back(sql_injection);
    
    // 大UDP包检测规则
    Rule large_udp;
    large_udp.name = "LARGE_UDP_PACKET";
    large_udp.description = "检测异常大的UDP数据包（可能是DoS攻击）";
    large_udp.src_ip = "*";
    large_udp.dst_ip = "*";
    large_udp.src_port = 0;
    large_udp.dst_port = 0;
    large_udp.protocol = UDP;
    large_udp.payload_pattern = "";
    large_udp.alert_on_match = true;
    large_udp.threshold = LARGE_UDP_SIZE; // 超过指定大小则告警
    detection_rules.push_back(large_udp);
    
    // SYN泛洪检测规则
    Rule syn_flood;
    syn_flood.name = "SYN_FLOOD";
    syn_flood.description = "检测TCP SYN泛洪攻击";
    syn_flood.src_ip = "*";
    syn_flood.dst_ip = "*";
    syn_flood.src_port = 0;
    syn_flood.dst_port = 0;
    syn_flood.protocol = TCP;
    syn_flood.payload_pattern = "";
    syn_flood.alert_on_match = true;
    syn_flood.threshold = SYN_FLOOD_THRESHOLD; // 每秒超过此数量则告警
    detection_rules.push_back(syn_flood);
}

// 计算最近n秒内的平均速率
int calculate_rate(const std::unordered_map<time_t, int>& rate_map, int seconds) {
    if (rate_map.empty()) return 0;
    
    time_t now = time(nullptr);
    time_t cutoff = now - seconds;
    int total = 0;
    
    for (const auto& entry : rate_map) {
        if (entry.first >= cutoff) {
            total += entry.second;
        }
    }
    
    return total / std::max(1, seconds); // 避免除以0
}

// 检测SYN泛洪攻击
bool detect_syn_flood(const std::string& dst_ip, const Session& session) {
    // 条件1: SYN包数量远大于ACK包
    if (session.syn_count > 0 && session.ack_count == 0 && session.syn_count > 5) {
        // 条件2: 计算最近10秒的SYN包速率
        int syn_rate = calculate_rate(session.packet_rate, 10);
        if (syn_rate > SYN_FLOOD_THRESHOLD) {
            return true;
        }
    }
    
    // 从目标角度检测: 来自多个源IP的SYN包
    if (target_stats.find(dst_ip) != target_stats.end()) {
        const TargetStats& stats = target_stats[dst_ip];
        // 源IP数量多但每个源发送的包少，典型的分布式攻击特征
        if (stats.src_ips.size() > 10 && stats.total_packets / stats.src_ips.size() < 5) {
            int tcp_rate = calculate_rate(stats.packet_rate, 10);
            if (tcp_rate > SYN_FLOOD_THRESHOLD * 2) {
                return true;
            }
        }
    }
    
    return false;
}

// 检测UDP泛洪攻击
bool detect_udp_flood(const std::string& dst_ip) {
    if (target_stats.find(dst_ip) == target_stats.end()) {
        return false;
    }
    
    const TargetStats& stats = target_stats[dst_ip];
    
    // 条件1: UDP包占比超过90% - 使用安全的常量映射访问方式
    int udp_count = get_map_value(stats.proto_counts, UDP, 0);
    if (udp_count * 10 < stats.total_packets * 9) { // udp_count < 90% of total
        return false;
    }
    
    // 条件2: 计算最近10秒的UDP速率
    int udp_rate = calculate_rate(stats.packet_rate, 10);
    if (udp_rate > UDP_FLOOD_THRESHOLD) {
        // 条件3: 包大小异常或payload无意义
        int avg_packet_size = stats.total_bytes / std::max(1, stats.total_packets);
        if (avg_packet_size > LARGE_UDP_SIZE || avg_packet_size < 10) {
            return true;
        }
    }
    
    return false;
}

// 检测ICMP泛洪攻击
bool detect_icmp_flood(const std::string& dst_ip) {
    if (target_stats.find(dst_ip) == target_stats.end()) {
        return false;
    }
    
    const TargetStats& stats = target_stats[dst_ip];
    
    // 条件1: ICMP包占比超过80% - 使用安全的常量映射访问方式
    int icmp_count = get_map_value(stats.proto_counts, ICMP, 0);
    if (icmp_count * 10 < stats.total_packets * 8) { // icmp_count < 80% of total
        return false;
    }
    
    // 条件2: 计算最近10秒的ICMP速率
    int icmp_rate = calculate_rate(stats.packet_rate, 10);
    if (icmp_rate > ICMP_FLOOD_THRESHOLD) {
        // 条件3: 多为Echo Request (Ping)
        return true;
    }
    
    return false;
}

// 检测HTTP泛洪攻击
bool detect_http_flood(const std::string& dst_ip, uint16_t dst_port) {
    if (dst_port != 80 && dst_port != 443) {
        return false;
    }
    
    if (target_stats.find(dst_ip) == target_stats.end()) {
        return false;
    }
    
    const TargetStats& stats = target_stats[dst_ip];
    
    // 计算最近10秒的HTTP请求速率
    int http_rate = calculate_rate(stats.packet_rate, 10);
    if (http_rate > HTTP_FLOOD_THRESHOLD) {
        // 条件1: 来源IP分布异常
        if (stats.src_ips.size() > 20 || stats.src_ips.size() == 1) {
            // 条件2: 包大小相似（可能是相同请求）
            int avg_packet_size = stats.total_bytes / std::max(1, stats.total_packets);
            int size_variation = 0;
            if (avg_packet_size > 0) {
                size_variation = 10; // 假设10%的变异系数
            }
            
            return true;
        }
    }
    
    return false;
}

// 检测Slowloris攻击
bool detect_slowloris(const std::string& src_ip, const std::string& dst_ip, 
                     uint16_t dst_port, const Session& session) {
    if (dst_port != 80 && dst_port != 443) {
        return false;
    }
    
    // 条件1: 连接建立时间长但数据量小
    time_t duration = session.last_seen - session.first_seen;
    if (duration > SLOWLORIS_TIMEOUT && session.total_bytes < 100) {
        // 条件2: 来自同一源IP的多个类似连接
        int similar_connections = 0;
        for (const auto& entry : sessions) {
            const Session& s = entry.second;
            if (s.src_ip == src_ip && s.dst_ip == dst_ip && 
                (s.dst_port == 80 || s.dst_port == 443) &&
                (s.last_seen - s.first_seen > SLOWLORIS_TIMEOUT) &&
                s.total_bytes < 100) {
                similar_connections++;
            }
        }
        
        if (similar_connections >= SLOWLORIS_THRESHOLD) {
            return true;
        }
    }
    
    return false;
}

// 检测IP分片攻击
bool detect_fragmentation_attack(const std::string& src_ip, const struct ip* ip_header) {
    // 条件1: 是分片包
    if ((ip_header->ip_off & htons(IP_MF | IP_OFFMASK)) == 0) {
        return false;
    }
    
    // 条件2: 统计来自该IP的分片包数量
    int fragment_count = 0;
    for (const auto& entry : sessions) {
        const Session& s = entry.second;
        if (s.src_ip == src_ip) {
            fragment_count += s.packet_count;
        }
    }
    
    if (fragment_count > FRAGMENT_THRESHOLD) {
        // 条件3: 分片偏移异常（如重叠）
        return true;
    }
    
    return false;
}

// 综合检测DoS攻击
AttackType detect_dos_attack(const std::string& src_ip, const std::string& dst_ip,
                            uint16_t src_port __attribute__((unused)), uint16_t dst_port, Protocol protocol,
                            const struct ip* ip_header, const Session& session) {
    // 按优先级检测不同类型的攻击
    if (protocol == TCP) {
        if (detect_slowloris(src_ip, dst_ip, dst_port, session)) {
            return SLOWLORIS;
        }
        if (detect_syn_flood(dst_ip, session)) {
            return SYN_FLOOD;
        }
        if (dst_port == 80 || dst_port == 443) {
            if (detect_http_flood(dst_ip, dst_port)) {
                return HTTP_FLOOD;
            }
        }
    } else if (protocol == UDP) {
        if (detect_udp_flood(dst_ip)) {
            return UDP_FLOOD;
        }
    } else if (protocol == ICMP) {
        if (detect_icmp_flood(dst_ip)) {
            return ICMP_FLOOD;
        }
    }
    
    // 检测分片攻击（适用于所有协议）
    if (detect_fragmentation_attack(src_ip, ip_header)) {
        return FRAGMENTATION_ATTACK;
    }
    
    return NONE;
}

// 检查规则匹配
void check_rules(const std::string& src_ip, const std::string& dst_ip,
                uint16_t src_port, uint16_t dst_port, Protocol protocol,
                const u_char* payload, int payload_len, const struct ip* ip_header,
                const Session& session) {
    // 先检查DoS攻击
    AttackType attack = detect_dos_attack(src_ip, dst_ip, src_port, dst_port, protocol, ip_header, session);
    if (attack != NONE) {
        std::string attack_name, attack_desc;
        switch (attack) {
            case SYN_FLOOD:
                attack_name = "SYN_FLOOD_ATTACK";
                attack_desc = "TCP SYN泛洪攻击 - 大量未完成的三次握手消耗目标资源";
                break;
            case UDP_FLOOD:
                attack_name = "UDP_FLOOD_ATTACK";
                attack_desc = "UDP泛洪攻击 - 大量UDP数据包消耗目标带宽和处理能力";
                break;
            case ICMP_FLOOD:
                attack_name = "ICMP_FLOOD_ATTACK";
                attack_desc = "ICMP泛洪攻击 - 大量Ping请求消耗目标资源";
                break;
            case HTTP_FLOOD:
                attack_name = "HTTP_FLOOD_ATTACK";
                attack_desc = "HTTP泛洪攻击 - 大量HTTP请求消耗Web服务器资源";
                break;
            case SLOWLORIS:
                attack_name = "SLOWLORIS_ATTACK";
                attack_desc = "Slowloris攻击 - 保持大量不完整HTTP连接耗尽服务器连接池";
                break;
            case FRAGMENTATION_ATTACK:
                attack_name = "FRAGMENTATION_ATTACK";
                attack_desc = "IP分片攻击 - 大量畸形分片包导致目标重组失败";
                break;
            default:
                return;
        }
        
        std::cout << "\033[1;31m[!] 严重告警: " << attack_name << "\033[0m" << std::endl;
        std::cout << "   描述: " << attack_desc << std::endl;
        std::cout << "   源: " << src_ip << ":" << src_port << std::endl;
        std::cout << "   目的: " << dst_ip << ":" << dst_port << std::endl;
        std::cout << "   协议: " << (protocol == TCP ? "TCP" : (protocol == UDP ? "UDP" : "ICMP")) << std::endl;
        
        // 输出攻击特征详情
        if (attack == SYN_FLOOD) {
            int syn_rate = calculate_rate(session.packet_rate, 10);
            std::cout << "   特征: SYN包速率 " << syn_rate << "个/秒，远超阈值" << SYN_FLOOD_THRESHOLD << std::endl;
        } else if (attack == UDP_FLOOD) {
            int udp_rate = calculate_rate(target_stats[dst_ip].packet_rate, 10);
            std::cout << "   特征: UDP包速率 " << udp_rate << "个/秒，远超阈值" << UDP_FLOOD_THRESHOLD << std::endl;
        }
        
        std::cout << "-------------------------" << std::endl;
    }
    
    // 检查其他规则
    for (const auto& rule : detection_rules) {
        bool match = true;
        
        if (rule.src_ip != "*" && rule.src_ip != src_ip)
            match = false;
        if (rule.dst_ip != "*" && rule.dst_ip != dst_ip)
            match = false;
        
        if (rule.src_port != 0 && rule.src_port != src_port)
            match = false;
        if (rule.dst_port != 0 && rule.dst_port != dst_port)
            match = false;
        
        if (rule.protocol != 0 && rule.protocol != protocol)
            match = false;
        
        if (!rule.payload_pattern.empty() && payload_len > 0) {
            std::string payload_str(reinterpret_cast<const char*>(payload), payload_len);
            if (payload_str.find(rule.payload_pattern) == std::string::npos)
                match = false;
        }
        
        // 规则特定的阈值检查
        if (rule.name == "LARGE_UDP_PACKET" && protocol == UDP) {
            if (payload_len < rule.threshold)
                match = false;
        }
        
        if (rule.name == "SSH_BRUTE_FORCE" && protocol == TCP && dst_port == 22) {
            if (session.packet_count < rule.threshold)
                match = false;
        }
        
        if (match && rule.alert_on_match) {
            std::cout << "\033[1;33m[!] 告警: " << rule.name << "\033[0m" << std::endl;
            std::cout << "   描述: " << rule.description << std::endl;
            std::cout << "   源: " << src_ip << ":" << src_port << std::endl;
            std::cout << "   目的: " << dst_ip << ":" << dst_port << std::endl;
            std::cout << "   协议: " << (protocol == TCP ? "TCP" : (protocol == UDP ? "UDP" : "未知")) << std::endl;
            std::cout << "-------------------------" << std::endl;
        }
    }
}

// 检查TCP异常
void check_tcp_anomalies(const tcphdr* tcp_header, const std::string& src_ip, 
                         const std::string& dst_ip, uint16_t src_port, uint16_t dst_port) {
    // 异常的TCP标志组合
    if ((tcp_header->syn == 1 && tcp_header->fin == 1) || 
        (tcp_header->syn == 1 && tcp_header->rst == 1) ||
        (tcp_header->fin == 1 && tcp_header->rst == 1)) {
        std::cout << "\033[1;33m[!] 告警: 异常的TCP标志组合\033[0m" << std::endl;
        std::cout << "   描述: 同时设置了不应该共存的TCP标志位，可能是扫描或攻击" << std::endl;
        std::cout << "   标志: SYN=" << (tcp_header->syn ? "1" : "0") << 
                     ", FIN=" << (tcp_header->fin ? "1" : "0") <<
                     ", RST=" << (tcp_header->rst ? "1" : "0") <<
                     ", ACK=" << (tcp_header->ack ? "1" : "0") << std::endl;
        std::cout << "   源: " << src_ip << ":" << src_port << std::endl;
        std::cout << "   目的: " << dst_ip << ":" << dst_port << std::endl;
        std::cout << "-------------------------" << std::endl;
    }
    
    // 异常的窗口大小
    if (ntohs(tcp_header->window) == 0 && tcp_header->ack == 1 && tcp_header->fin == 0) {
        std::cout << "\033[1;33m[!] 告警: 异常的TCP窗口大小\033[0m" << std::endl;
        std::cout << "   描述: 窗口大小为0但未发送FIN，可能是流量控制异常" << std::endl;
        std::cout << "   源: " << src_ip << ":" << src_port << std::endl;
        std::cout << "   目的: " << dst_ip << ":" << dst_port << std::endl;
        std::cout << "-------------------------" << std::endl;
    }
}

// 数据包处理回调函数
void packet_handler(u_char *user_data __attribute__((unused)), 
                   const struct pcap_pkthdr *pkthdr, 
                   const u_char *packet) {
    packet_counter++;
    
    // 以太网头部解析
    struct ethhdr *eth_header = (struct ethhdr *)packet;
    
    // 只处理IPv4数据包
    if (ntohs(eth_header->h_proto) != 0x0800) {
        return;
    }
    
    // IP头部解析
    const struct ip *ip_header = (struct ip*)(packet + sizeof(struct ethhdr));
    int ip_header_len = ip_header->ip_hl * 4;
    
    // 源IP和目的IP
    std::string src_ip = inet_ntoa(ip_header->ip_src);
    std::string dst_ip = inet_ntoa(ip_header->ip_dst);
    
    // 协议相关变量
    Protocol protocol = UNKNOWN;
    uint16_t src_port = 0, dst_port = 0;
    const u_char* payload = nullptr;
    int payload_len = 0;
    const tcphdr* tcp_header = nullptr;
    const udphdr* udp_header = nullptr;
    const icmphdr* icmp_header = nullptr;
    
    // 解析传输层协议
    switch (ip_header->ip_p) {
        case IPPROTO_TCP: 
        {
            protocol = TCP;
            tcp_header = (struct tcphdr*)(packet + sizeof(struct ethhdr) + ip_header_len);
            src_port = ntohs(tcp_header->source);
            dst_port = ntohs(tcp_header->dest);
            int tcp_header_len = tcp_header->doff * 4;
            payload = packet + sizeof(struct ethhdr) + ip_header_len + tcp_header_len;
            payload_len = pkthdr->len - (sizeof(struct ethhdr) + ip_header_len + tcp_header_len);
            
            check_tcp_anomalies(tcp_header, src_ip, dst_ip, src_port, dst_port);
            break;
        }
            
        case IPPROTO_UDP: 
        {
            protocol = UDP;
            udp_header = (struct udphdr*)(packet + sizeof(struct ethhdr) + ip_header_len);
            src_port = ntohs(udp_header->source);
            dst_port = ntohs(udp_header->dest);
            payload = packet + sizeof(struct ethhdr) + ip_header_len + sizeof(struct udphdr);
            payload_len = ntohs(udp_header->len) - sizeof(struct udphdr);
            if (payload_len < 0) payload_len = 0;
            break;
        }
            
        case IPPROTO_ICMP: 
        {
            protocol = ICMP;
            icmp_header = (struct icmphdr*)(packet + sizeof(struct ethhdr) + ip_header_len);
            payload = packet + sizeof(struct ethhdr) + ip_header_len + sizeof(struct icmphdr);
            payload_len = pkthdr->len - (sizeof(struct ethhdr) + ip_header_len + sizeof(struct icmphdr));
            break;
        }
            
        default:
            return;
    }
    
    // 更新会话和目标统计
    update_session(src_ip, dst_ip, src_port, dst_port, protocol, payload, payload_len, tcp_header);
    update_target_stats(src_ip, dst_ip, protocol, pkthdr->len);
    
    // 获取当前会话信息用于检测
    std::string session_id = generate_session_id(src_ip, dst_ip, src_port, dst_port, protocol);
    Session& current_session = sessions[session_id];
    
    // 检查规则和攻击
    check_rules(src_ip, dst_ip, src_port, dst_port, protocol, payload, payload_len, ip_header, current_session);
    
    // 定期清理过期会话（每1000个包执行一次）
    if (packet_counter % 1000 == 0) {
        time_t now = time(nullptr);
        for (auto it = sessions.begin(); it != sessions.end();) {
            if (now - it->second.last_seen > 300) { // 5分钟无活动
                it = sessions.erase(it);
            } else {
                ++it;
            }
        }
        
        std::cout << "\033[1;34m[*] 已处理 " << packet_counter << " 个数据包，当前活跃会话: " << sessions.size() << "\033[0m" << std::endl;
    }
}

// 显示可用网络接口并返回接口数量
int list_interfaces() {
    pcap_if_t *all_devs;
    pcap_if_t *d;
    int i = 0;
    char errbuf[PCAP_ERRBUF_SIZE];
    
    if (pcap_findalldevs(&all_devs, errbuf) == -1) {
        std::cerr << "获取接口列表失败: " << errbuf << std::endl;
        exit(1);
    }
    
    std::cout << "可用网络接口:" << std::endl;
    for (d = all_devs, i = 0; d != nullptr; d = d->next, i++) {
        std::cout << i << ". " << d->name << std::endl;
        if (d->description)
            std::cout << "   描述: " << d->description << std::endl;
        else
            std::cout << "   无描述信息" << std::endl;
    }
    
    if (i == 0) {
        std::cerr << "未找到任何网络接口! 请检查权限。" << std::endl;
        exit(1);
    }
    
    pcap_freealldevs(all_devs);
    return i;  // 返回接口数量
}

int main(int argc, char *argv[]) {
    program_start_time = time(nullptr);
    
    char errbuf[PCAP_ERRBUF_SIZE];
    pcap_if_t *all_devs;
    pcap_if_t *d;
    pcap_t *handle;
    int inum, i, interface_count;
    char filter_exp[1024] = "";
    struct bpf_program fp;
    bpf_u_int32 net, mask;
    
    std::cout << "=== 网络协议安全分析与漏洞检测工具 ===" << std::endl;
    
    load_rules();
    std::cout << "[*] 已加载 " << detection_rules.size() << " 条检测规则" << std::endl;
    std::cout << "[*] DoS攻击检测已启用，包含SYN泛洪、UDP泛洪、ICMP泛洪等类型" << std::endl;
    
    interface_count = list_interfaces();
    
    std::cout << "请输入要监听的接口编号 (0-" << (interface_count - 1) << "): ";
    std::cin >> inum;
    
    if (inum < 0 || inum >= interface_count) {
        std::cerr << "无效的接口编号! 请输入0到" << (interface_count - 1) << "之间的数字。" << std::endl;
        exit(1);
    }
    
    if (pcap_findalldevs(&all_devs, errbuf) == -1) {
        std::cerr << "获取接口列表失败: " << errbuf << std::endl;
        exit(1);
    }
    
    for (d = all_devs, i = 0; i < inum; d = d->next, i++);
    
    if (argc > 1) {
        strncpy(filter_exp, argv[1], sizeof(filter_exp) - 1);
        std::cout << "[*] 使用过滤规则: " << filter_exp << std::endl;
    }
    
    handle = pcap_open_live(d->name, BUFSIZ, 1, 1000, errbuf);
    if (handle == nullptr) {
        std::cerr << "无法打开设备 " << d->name << ": " << errbuf << std::endl;
        pcap_freealldevs(all_devs);
        exit(1);
    }
    
    if (pcap_datalink(handle) != DLT_EN10MB) {
        std::cerr << d->name << " 不是以太网接口，不支持" << std::endl;
        pcap_freealldevs(all_devs);
        exit(1);
    }
    
    if (pcap_lookupnet(d->name, &net, &mask, errbuf) == -1) {
        std::cerr << "无法获取网络号和子网掩码，使用默认值: " << errbuf << std::endl;
        net = 0;
        mask = 0;
    }
    
    if (filter_exp[0] != '\0') {
        if (pcap_compile(handle, &fp, filter_exp, 0, net) == -1) {
            std::cerr << "过滤规则编译失败: " << pcap_geterr(handle) << std::endl;
            pcap_freealldevs(all_devs);
            exit(1);
        }
        if (pcap_setfilter(handle, &fp) == -1) {
            std::cerr << "设置过滤规则失败: " << pcap_geterr(handle) << std::endl;
            pcap_freealldevs(all_devs);
            exit(1);
        }
    }
    
    std::cout << "[*] 开始监听接口 " << d->name << " ..." << std::endl;
    std::cout << "[*] 按 Ctrl+C 停止" << std::endl;
    
    pcap_freealldevs(all_devs);
    pcap_loop(handle, 0, packet_handler, nullptr);
    pcap_close(handle);
    
    std::cout << "\n[*] 已停止，共处理 " << packet_counter << " 个数据包" << std::endl;
    
    return 0;
}
     
    
